from mitmproxy import proxy, options, controller, master
from mitmproxy.proxy.server import ProxyServer
from mitmproxy import http
from .models import proxy_data
from lib.config_json import load_config
from celery.decorators import task
import os

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))




PROXY_LISTEN_PORT = load_config('PROXY_LISTEN_PORT')
PROXY_LISTEN_HOST = load_config('PROXY_LISTEN_HOST')
PROXY_LISTEN_MODE = load_config('PROXY_LISTEN_MODE')
EXCLUDE_STATIC_FILE = load_config('EXCLUDE_STATIC_FILE')


class ProxyMaster(master.Master):
    def __init__(self, opts, server):
        super(ProxyMaster, self).__init__(opts, server)

    def run(self):
        # try:
        print("开始代理")
        master.Master.run(self)
        # except KeyboardInterrupt:
        #    self.shutdown()

    def stop(self):
        self.shutdown()

    @controller.handler
    def request(self, f):
        proxy_request_handle(f)

    @controller.handler
    def response(self, f):
        proxy_response_handle(f)
        # parser = ResponseParser(f)
        # insert_result(parser.parser_data())


def proxy_request_handle(flow: http.HTTPFlow):
    pass
    # change the request headers['Host']
    # flow.request.headers['X-Online-Host'] = 'wap.gd.10086.cn'


def proxy_response_handle(flow: http.HTTPFlow):
    path = '/{}'.format('/'.join(flow.request.path_components))
    try:
        client_ip = str(flow.client_conn.ip_address).split(':')[0]
    except:
        client_ip = '0.0.0.0'
    code = 0
    for i in EXCLUDE_STATIC_FILE:  # Exclude static file in capture packet
        if path.find(i) > 0:
            code = 1
            break
    if code == 0:
        headers = ""
        cookies = ""
        for i in flow.request.headers:
            headers = headers + i + ":" + flow.request.headers[i] + "\n"
        for i in flow.request.cookies:
            cookies = cookies + i+"="+flow.request.cookies[i]+";"
        proxy_data.objects.create(
            # save data in database
            host=flow.request.host+":"+str(flow.request.port),

            method=flow.request.method,
            url=flow.request.url,
            scheme=flow.request.scheme,
            request_headers=headers,
            request_content=str(flow.request.content, encoding="utf8"),
            path=path,
            status_code=flow.response.status_code,
            response_headers=flow.response.headers,
            response_content=flow.response.content,
            request_cookies=cookies,
            http_version=flow.request.http_version,
            ip = client_ip,
        )



def start_proxy():
    opts = options.Options(
        listen_host=PROXY_LISTEN_HOST,
        listen_port=PROXY_LISTEN_PORT,
        mode=PROXY_LISTEN_MODE,
        cadir=BASE_DIR+"/ssl/",
    )

    config = proxy.ProxyConfig(opts)
    try:
        server = ProxyServer(config)

        m = ProxyMaster(opts, server)
        m.run()
    except:
        pass


# 不能停止代理，待研究
def shutdown_proxy():
    opts = options.Options(
        listen_host=PROXY_LISTEN_HOST,
        listen_port=PROXY_LISTEN_PORT,
        mode=PROXY_LISTEN_MODE,
        cadir=PROXY_CACER_DIR,
    )

    config = proxy.ProxyConfig(opts)
    server = ProxyServer(config)

    m = master.Master(opts, server)
    m.shutdown()
